using namespace std;
#include<iostream>
#include<iomanip>
#include<fstream>
#include<vector>
#include<stdio.h>
#include<stdlib.h>
#include<map>
#include<set>
#include"BA2.h"
#include<math.h>
#include<gsl/gsl_rng.h>
#include<gsl/gsl_randist.h>
#include<gsl/gsl_sf_gamma.h>

extern std::vector<string> cmdList;

// parameters for simulation program
double fst=0.05;
double avgMigration=0.05;
long int numPoplns=2;
long int numLoci=10;
long int sampleSize=50;
vector<double> Fstat;
string simfilename="simulate.out";


// constructor for class Simulate

Simulate::Simulate()
{
  vector<double> tempLocus;
  for(int i=0;i<numPoplns;i++)
    sampleSizes.push_back(sampleSize);
  if(Fstat.size()==0)
    for(int i=0;i<numPoplns*sampleSize;i++)
      Fstat.push_back(0.0);

  simSeed=1;
  for(int i=0;i<=2;i++)
    {
      tempLocus.push_back(1.0/3.0);
    }
  for(int i=0;i<numLoci;i++)
    {
      avgAlleleFreqs.push_back(tempLocus);
    }
  for(int i=0;i<numPoplns;i++)
    {
      vector<double> tempMig;
      for(int j=0;j<numPoplns;j++)
	{
	  if(i!=j)
	    tempMig.push_back(avgMigration/numPoplns);
	  else
	    tempMig.push_back(1.0-avgMigration*3.0);
	}
      simMigrationRates.push_back(tempMig);
    }
}

void Simulate::simulateSample()
{ 
  ofstream simout;
  simout.open(simfilename.c_str());
  if(simout)
    {
      for(int j=0;j<numPoplns;j++)
	{
	  vector<int> temp;
	  for(int i=0;i<sampleSizes[j];i++)
	    {
	      temp.push_back(0);
	    }
	  migrantAncestry.push_back(temp);
	  ancestryAge.push_back(temp);
	}
      
      // initiate random number generator
      gsl_rng *r;	
      r=gsl_rng_alloc(gsl_rng_taus);
      gsl_rng_set (r,simSeed); 
      
      double tempAF[500], tempAP[500];
      
      // simulate allele frequencies
      for(int j=0;j<numPoplns;j++)
	{
	  vector<vector<double> > tempPoplns;
	  for(int k=0;k<numLoci;k++)
	    {
	      for(int i=0;i<static_cast<int>(avgAlleleFreqs[k].size());i++)
		{
		  tempAP[i]=avgAlleleFreqs[k][i]*(1.0/fst-1.0);
		}
	      gsl_ran_dirichlet(r,avgAlleleFreqs[k].size(),tempAP,tempAF);
	      vector<double> tempLocus;
	      for(int i=0;i<static_cast<int>(avgAlleleFreqs[k].size());i++)
		tempLocus.push_back(tempAF[i]);
	      tempPoplns.push_back(tempLocus);
	    }
	  simAlleleFreqs.push_back(tempPoplns);
	}
      
      // simulate migrant ancestries
      for(int i=0;i<numPoplns;i++)
	for(int j=0;j<sampleSizes[i];j++)
	  {
	    double runi=gsl_rng_uniform_pos(r);
	    double migRSum=0.0;
	    int migrantPop=0;
	    while ((runi>migRSum)&&(migrantPop<numPoplns))
	      {
		if(migrantPop!=i)
		  migRSum+=3.0*simMigrationRates[migrantPop][i];
		++migrantPop;
	      }
	    --migrantPop;
	    if(runi<migRSum)
	      migrantAncestry[i][j]=migrantPop;
	    else
	      migrantAncestry[i][j]=i;
	    if(migrantAncestry[i][j]==i)
	      ancestryAge[i][j]=0;
	    else
	      if(gsl_rng_uniform_pos(r)<=0.333)
		ancestryAge[i][j]=1;
	      else
		ancestryAge[i][j]=2;
	    
	  }
      
      // simulate genotypes
      int indivIndex=0;
      for(int i=0;i<numPoplns;i++)
	{
	  for(int j=0;j<sampleSizes[i];j++)
	    {
	      simIndiv tempIndiv;
	      simGene geno;
	      if(ancestryAge[i][j]==0)
		{
		  for(int k=0;k<numLoci;k++)
		    {
		      if(Fstat[indivIndex]>0)
			{
			  double runi=gsl_rng_uniform_pos(r);
			  if(runi<=Fstat[indivIndex])
			    {
			      runi=gsl_rng_uniform_pos(r);
			      double sumFreq=0.0;
			      int iter1=0;
			      while((runi>sumFreq)&&(iter1<static_cast<int>(simAlleleFreqs[i][k].size())))
				{
				  sumFreq+=simAlleleFreqs[i][k][iter1];
				  ++iter1;
				}
			      geno.g1=iter1-1;
			      geno.g2=iter1-1;
			      tempIndiv.genotypes.push_back(geno);
			    }
			  else
			    {
			      double runi=gsl_rng_uniform_pos(r);
			      double sumFreq=0.0;
			      int iter1=0;
			      while((runi>sumFreq)&&(iter1<static_cast<int>(simAlleleFreqs[i][k].size())))
				{
				  sumFreq+=simAlleleFreqs[i][k][iter1];
				  ++iter1;
				}
			      geno.g1=iter1-1;
			      runi=gsl_rng_uniform_pos(r);
			      sumFreq=0.0;
			      iter1=0;
			      while((runi>sumFreq)&&(iter1<static_cast<int>(simAlleleFreqs[i][k].size())))
				{
				  sumFreq+=simAlleleFreqs[i][k][iter1];
				  ++iter1;
				}
			      geno.g2=iter1-1;
			      tempIndiv.genotypes.push_back(geno);
			    }			  
			}
		      else
			{
			  double runi=gsl_rng_uniform_pos(r);
			  double sumFreq=0.0;
			  int iter1=0;
			  while((runi>sumFreq)&&(iter1<static_cast<int>(simAlleleFreqs[i][k].size())))
			    {
			      sumFreq+=simAlleleFreqs[i][k][iter1];
			      ++iter1;
			    }
			  geno.g1=iter1-1;
			  runi=gsl_rng_uniform_pos(r);
			  sumFreq=0.0;
			  iter1=0;
			  while((runi>sumFreq)&&(iter1<static_cast<int>(simAlleleFreqs[i][k].size())))
			    {
			      sumFreq+=simAlleleFreqs[i][k][iter1];
			      ++iter1;
			    }
			  geno.g2=iter1-1;
			  tempIndiv.genotypes.push_back(geno);
			}
		    }
		}
	      else
		if(ancestryAge[i][j]==1)
		  {
		    for(int k=0;k<numLoci;k++)
		      {
			if(Fstat[indivIndex]>0)
			  {
			    double runi=gsl_rng_uniform_pos(r);
			    if(runi<=Fstat[indivIndex])
			      {
				runi=gsl_rng_uniform_pos(r);
				double sumFreq=0.0;
				int iter1=0;
				while((runi>sumFreq)&&(iter1<static_cast<int>(simAlleleFreqs[i][k].size())))
				  {
				    sumFreq+=simAlleleFreqs[migrantAncestry[i][j]][k][iter1];
				    ++iter1;
				  }
				geno.g1=iter1-1;
				geno.g2=iter1-1;
				tempIndiv.genotypes.push_back(geno);
			      }
			    else
			      {
				double runi=gsl_rng_uniform_pos(r);
				double sumFreq=0.0;
				int iter1=0;
				while((runi>sumFreq)&&(iter1<static_cast<int>(simAlleleFreqs[i][k].size())))
				  {
				    sumFreq+=simAlleleFreqs[migrantAncestry[i][j]][k][iter1];
				    ++iter1;
				  }
				geno.g1=iter1-1;
				runi=gsl_rng_uniform_pos(r);
				sumFreq=0.0;
				iter1=0;
			      while((runi>sumFreq)&&(iter1<static_cast<int>(simAlleleFreqs[i][k].size())))
				{
				  sumFreq+=simAlleleFreqs[migrantAncestry[i][j]][k][iter1];
				  ++iter1;
				}
			      geno.g2=iter1-1;
			      tempIndiv.genotypes.push_back(geno);
			    }			  
			}
			else
			  {
			    double runi=gsl_rng_uniform_pos(r);
			    double sumFreq=0.0;
			    int iter1=0;
			    while((runi>sumFreq)&&(iter1<static_cast<int>(simAlleleFreqs[i][k].size())))
			      {
				sumFreq+=simAlleleFreqs[migrantAncestry[i][j]][k][iter1];
				++iter1;
			      }
			    geno.g1=iter1-1;
			    runi=gsl_rng_uniform_pos(r);
			    sumFreq=0.0;
			    iter1=0;
			    while((runi>sumFreq)&&(iter1<static_cast<int>(simAlleleFreqs[i][k].size())))
			      {
				sumFreq+=simAlleleFreqs[migrantAncestry[i][j]][k][iter1];
				++iter1;
			      }
			    geno.g2=iter1-1;
			    tempIndiv.genotypes.push_back(geno);
			  }
		      }      
		  }
		else
		  if(ancestryAge[i][j]==2)
		    {
		      for(int k=0;k<numLoci;k++)
			{
			  double runi=gsl_rng_uniform_pos(r);
			  double sumFreq=0.0;
			  int iter1=0;
			  while((runi>sumFreq)&&(iter1<static_cast<int>(simAlleleFreqs[i][k].size())))
			    {
			      sumFreq+=simAlleleFreqs[migrantAncestry[i][j]][k][iter1];
			      ++iter1;
			    }
			  geno.g1=iter1-1;
			  runi=gsl_rng_uniform_pos(r);
			  sumFreq=0.0;
			  iter1=0;
			  while((runi>sumFreq)&&(iter1<static_cast<int>(simAlleleFreqs[i][k].size())))
			    {
			      sumFreq+=simAlleleFreqs[i][k][iter1];
			      ++iter1;
			    }
			  geno.g2=iter1-1;
			  tempIndiv.genotypes.push_back(geno);
			}
		    }
	      tempIndiv.sampledPopulation=i;
	      tempIndiv.label=indivIndex;
	      simulatedSample.push_back(tempIndiv);
	      ++indivIndex;
	    }
	}

      // create input file from simulated data
      for(int i=0;i<static_cast<int>(simulatedSample.size());i++)
	{
	  simout << simulatedSample[i].label << " " << simulatedSample[i].sampledPopulation << " ";
	  for(int j=0;j<static_cast<int>(simulatedSample[i].genotypes.size());j++)
	    simout << simulatedSample[i].genotypes[j].g1 << " ";
	  simout << "\n";
	  simout << simulatedSample[i].label << " " << simulatedSample[i].sampledPopulation << " ";
	  for(int j=0;j<static_cast<int>(simulatedSample[i].genotypes.size());j++)
	    simout << simulatedSample[i].genotypes[j].g2 << " ";
	  simout << "\n";
	}
      simout.close();
    }
  else
    {
      cout << "Failed to open simulation output file: " << simfilename << "\n";
      exit(0);
    }
}

void PrintSimulationParameters()
{
	  cout << "                                                            \n";
	  cout << "------------------------------------------------------------\n";
	  cout << "               BayesAss V.2.0: Simulation Settings           \n";
	  cout << "------------------------------------------------------------\n";
	  cout << " fst=" << fst << "\n";
	  cout << " avgMigration=" << avgMigration << "\n";
	  cout << " sampleSize=" << sampleSize << "\n";
	  cout << " numLoci=" << numLoci << "\n";
	  cout << " numPoplns=" << numPoplns << "\n";
	  cout << " simfile=" << simfilename << "\n";
	  cout << " Fstat=(";
	  for(int i=0;i<numPoplns*sampleSize-1;i++)
	    cout << Fstat[i] << ",";
	  cout << Fstat[numPoplns*sampleSize-1] << ")\n";
	  cout << "------------------------------------------------------------\n";
	  cout << "                                                            \n";
}

void sset()
{
  string parm;
  string paramValue;
  if(cmdList.size()==1)
    {
      if(Fstat.size()==0)
	for(int i=0;i<numPoplns*sampleSize;i++)
	  Fstat.push_back(0.0);
      PrintSimulationParameters();
    }
  else
    if(cmdList.size()==2)
      {
	int pos=cmdList[1].find_first_of("=");
	if(pos!=static_cast<int>(string::npos))
	  {
	    parm=cmdList[1].substr(0,pos);
	    paramValue=cmdList[1].substr(pos+1);
	  }
	if(static_cast<int>(Fstat.size())==0)
	  for(int i=0;i<numPoplns*sampleSize;i++)
	    Fstat.push_back(0.0);
	if(parm=="avgMigration") avgMigration=atof(paramValue.c_str());
	else
	  if(parm=="sampleSize") 
	    {
	      sampleSize=atol(paramValue.c_str());
	      if(static_cast<int>(Fstat.size())>(sampleSize*numPoplns))
		Fstat.erase(Fstat.begin()+sampleSize*numPoplns,Fstat.end());
	      else
		if(static_cast<int>(Fstat.size())<(sampleSize*numPoplns))
		  for(int i=0;i<(sampleSize*numPoplns)-static_cast<int>(Fstat.size());i++)
		    Fstat.push_back(0.0);
	    }
	  else
	    if(parm=="numLoci") numLoci=atol(paramValue.c_str());
	    else
	      if(parm=="fst") fst=atof(paramValue.c_str());
	      else
		if(parm=="simfile") simfilename=paramValue;
		else
		  if(parm=="numPoplns") 
		    {
		      numPoplns=atol(paramValue.c_str());
		      if(static_cast<int>(Fstat.size())>(sampleSize*numPoplns))
			Fstat.erase(Fstat.begin()+sampleSize*numPoplns,Fstat.end());
		      else
			if(static_cast<int>(Fstat.size())<(sampleSize*numPoplns))
			  for(int i=0;i<(sampleSize*numPoplns)-static_cast<int>(Fstat.size());i++)
			    Fstat.push_back(0.0);
		    }
		  else
		    if(parm=="Fstat")
		      {
			int pos=0, prev_pos=0;
			string temp1, temp2;
			int FstatPos;
			double FstatVal;
			paramValue.erase(paramValue.begin(),paramValue.begin()+1);
			if((pos=paramValue.find_first_of(',',pos))!=static_cast<int>(string::npos))
			  {
			    temp1=paramValue.substr(prev_pos,pos-prev_pos);
			    prev_pos = ++pos;
			  }
			if((pos=paramValue.find_first_of(')',pos))!=static_cast<int>(string::npos))
			  {
			    temp2=paramValue.substr(prev_pos,pos-prev_pos);
			    prev_pos = ++pos;
			  }
			FstatPos=atol(temp1.c_str());
			FstatVal=atof(temp2.c_str());
			if(FstatPos<static_cast<int>(Fstat.size()))
			  Fstat[FstatPos]=FstatVal;
		      }
		      
	PrintSimulationParameters();
      }
}

void Simulate::PrintAlleleFreqs()
{
  for(int i=0;i<numPoplns;i++)
    {
      cout << "\nPopulation: " << i << "\n";
      for(int j=0;j<static_cast<int>(simAlleleFreqs[i].size());j++)
	{
	  cout << "Locus " << j << ": ";
	  for(int k=0;k<static_cast<int>(simAlleleFreqs[i][j].size());k++)
	    cout << k << ":" << simAlleleFreqs[i][j][k] << " ";
	  cout << "\n";
	} 
	}
  cout << "\n";
}

void Simulate::PrintMigrantAncestries()
{
  int indivIn=0;
  for(int i=0;i<numPoplns;i++)
    for(int j=0;j<sampleSizes[i];j++)
      {
	cout << " popln:" << i << " indiv:" << indivIn << " migrantPop:";
	cout << migrantAncestry[i][j] << " migrantAge:" << ancestryAge[i][j] << "\n";
	++indivIn;
      }
}
